test(model-loaders): add regression tests for checkpoint key-remapping#9336
Open
Pfannkuchensack wants to merge 2 commits into
Open
test(model-loaders): add regression tests for checkpoint key-remapping#9336Pfannkuchensack wants to merge 2 commits into
Pfannkuchensack wants to merge 2 commits into
Conversation
Extract the state-dict key surgery in the qwen_image, flux, and anima loaders into pure module-level helpers and cover them with unit tests, so a regression like the transformers-5.x Qwen-VL remap bug fails in CI instead of at load time. - qwen_image: extract _remap_qwen_vl_checkpoint_keys - flux: move the FLUX.2 BFL->diffusers transformer/VAE converters into a new flux2_state_dict_utils module, deduplicating the two identical copies that lived in the checkpoint and GGUF loader classes - anima: extract _strip_anima_bundle_prefix Tests use fixtures captured from real single-file checkpoints and, where the diffusers architecture builds cheaply, assert on the meta device that every converted key maps to a real parameter (the FLUX.2 VAE is checked for full bijective coverage). Behavior of all loaders is unchanged.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR fixes a FLUX/Qwen single-file loading regression and adds the regression-test coverage that would have caught it.
The fix.
Qwen2_5_VLForConditionalGeneration._checkpoint_conversion_mappingis empty in transformers 5.x, so theif key_mapping:guard skipped remapping and ComfyUI single-file Qwen2.5-VL encoder checkpoints (visual.*,model.*) failed to load: their keys were reported as unexpected whilemodel.visual.*params stayed as meta tensors. We fall back to the legacy mapping when transformers no longer provides one; a negative lookahead keeps new-layout checkpoints untouched.The tests. Extract the state-dict key surgery in the qwen_image, flux, and anima loaders into pure module-level helpers and cover them with unit tests, so a regression like this one fails in CI instead of at load time.
_remap_qwen_vl_checkpoint_keysflux2_state_dict_utilsmodule, deduplicating the two identical copies that lived in the checkpoint and GGUF loader classes_strip_anima_bundle_prefixTests use fixtures captured from real single-file checkpoints and, where the diffusers architecture builds cheaply, assert on the meta device that every converted key maps to a real parameter (the FLUX.2 VAE is checked for full bijective coverage). Behavior of all loaders is unchanged.
Related Issues / Discussions
N/A
QA Instructions
Automated (no network, no GPU, no real weights):
uv run --extra cuda --extra test python -m pytest tests/backend/model_manager/load/ --no-covload/suite stays green (246 passed, 122 skipped)._remap_qwen_vl_checkpoint_keys→test_qwen_image_state_dict_utils.pygoes red (incl. the bug-report-keys test), proving the tests catch the original bug.Manual (optional): load a ComfyUI single-file Qwen2.5-VL encoder (e.g.
qwen_2.5_vl_7b_fp8_scaled.safetensors) and run aqwen_image_text_encoderinvocation — before the fix it raisedFailed to load all parameters from checkpoint. Meta tensors remain: [...]with a728 unexpected keyswarning; after the fix it loads cleanly.Merge Plan
Standard merge — no DB/schema changes, no migrations. Loader behavior is unchanged (converters were only extracted/deduplicated), validated end-to-end against real FLUX.2 transformer/VAE checkpoints (0 unmatched keys).
Checklist
What's Newcopy (if doing a release after this PR)